# Eisenstein.py : By Sean Starkweather 

import numpy as np
import matplotlib.pyplot as plt
import argparse
from sympy import * 

# eistenstein.py: 
# Generate a partial triangular tesselation of the complex plane,
# whose outer boundary is a hexagon, using the cube roots of unity,
# (1, w, w^2) where w = 1/2(-1 - i sqrt(3)
# Integers of the form, a+bw, are the Eisenstein-Jacobi integers and form a 
# field. Once the tesselation has been generated, then go search 
# for the primes in the field. The program gives an option for plotting them.  

w = 0.5*(-1+1.7320508075688772j)
z = symbols('z origin_sym', complex=True)
# z  symbolic form of w
z = Rational(1,2)*(-1 + sqrt(3)*I)

# origin
origin_sym = 0+0*I
origin = 0+0j

# hull_gen(n)
# Generate a hexagonal hull of points of radius n from the origin.
# returns array


# lattice_gen(n)
# Generate a triangular lattice of radius n, a sum over hexagonal hulls, 
# working outward from the origin. The number of points generated for a given n 
# are given by O(n) = 3n(n+1)+1, the proof of this identity follows from a 
# simple application of Gauss' expression for the sum of integers 
# from 1 to N, or (1/2)N(N+1). 
# returns array


# Plot an array of complex numbers in the complex plane
def complex_plot(a):
    for x in range(len(a)):
        plt.plot([0,a[x].real],[0,a[x].imag],'r.',label='python')
    limit=np.max(np.ceil(np.absolute(a))) # set limits for axis
    plt.xlim((-limit,limit))
    plt.ylim((-limit,limit))
    plt.ylabel('Imaginary')
    plt.xlabel('Real')
    plt.show()

def print_arr(arr):
    i=0
    for e in arr:
	print i, ': x: ',e.real,'\t  y:', e.imag
	i+=1

# generate eisenstein coordinates for array of integers
# returns coords



class Eisenstein_int:
	float_coord = []
	exact_coord = []
	eisenstein_coord = [0,0]
	modulus_squared = 0
	prime = False

	def iseprime(self):

		a = self.eisenstein_coord[0]
		b = self.eisenstein_coord[1]

		if ((b == 0) and isprime(abs(a))):
			if ((abs(a) % 3) == 2):
				return True

		if (isprime(self.modulus_squared) and (self.modulus_squared % 3 == 1)):
			return True

		if(isprime(self.modulus_squared) and (self.modulus_squared % 3 == 0)):
			return True
		
		if(isprime(sqrt(self.modulus_squared)) and (sqrt(self.modulus_squared) % 3 == 2)):
			return True
				
		return False

	def __init__(self, float_coord, exact_coord, eisenstein_coord):
		a = b = 0
		self.float_coord = float_coord
		self.exact_coord = exact_coord
		self.eisenstein_coord = eisenstein_coord
		a = eisenstein_coord[0]
		b = eisenstein_coord[1]
		self.modulus_squared = a**2 + b**2 - a*b
		self.prime = self.iseprime()
 	
class Eisenstein_mesh:
	points = []
	grid = []
	index = 0

	def __init__(self, n):
		grid = self.lattice_gen(n)
		coords = self.eisenstein_coords(grid[1])

		for i in range (0, len(grid[0])): 
			self.points.append(Eisenstein_int(grid[0][i], (grid[1][i]).expand(), coords[i]))
			self.index +=1

	def hull_gen(self, n):
		hull = rot = [] 
		hull_sym = rot_sym = []
	
		if (n==0):
			hull = np.append(hull, origin)
			hull_sym = np.append(hull_sym, origin_sym)
			return [hull, hull_sym]

		# generate first set of units
		for e in range(0, n+1):
			hull = np.append(hull, n+e*w)
			hull_sym = np.append(hull_sym, n+e*z)

		for e in range(1, n):
	 		hull = np.append(hull, (n-e) + n*w)
	 		hull_sym = np.append(hull_sym, (n-e) + n*z)

		# multiply units by w, a rotation of points by 2pi/3 radians (120 degrees)
		rot = w*hull
		hull = np.append(hull, rot)
	
		rot_sym = z*hull_sym
		hull_sym = np.append(hull_sym, rot_sym)

		# multiply again by w, to get the 'w^2' points and close the hull.
		rot = w*rot
		hull = np.append(hull, rot)

		rot_sym = z*rot_sym
		hull_sym = np.append(hull_sym, rot_sym)
		return [hull, hull_sym]
	
	def lattice_gen(self, n):
		lattice = lattice_sym = []
		lat_pair = []
		for i in range(0,n+1):
			lat_pair = self.hull_gen(i) 	
			lattice = np.append(lattice, lat_pair[0])	
			lattice_sym = np.append(lattice_sym, lat_pair[1])	
		pair = [lattice, lattice_sym]
		return pair

	def eisenstein_coords(self, arr):
		coords = []
		for e in arr:
			b = im(e)*Pow(sqrt(3), -1)*(-2)
			a = re(e)-im(e)*Pow(sqrt(3), -1)
			elt = [a,b]
			coords.append(elt)
		return coords

# main	
parser = argparse.ArgumentParser()
parser.add_argument("n", help='generate grid of order n', type=int )
parser.add_argument("-d", help='dump grid of order n to stdout', action="store_true")
parser.add_argument("-ds", help='dump grid of order n exact vectors to stdout', action="store_true")
parser.add_argument("-de", help='dump eisenstein coordinates for grid of order n', action="store_true")
parser.add_argument("-da", help='dump all exact, numeric and eisenstein coordinates for a grid of order n', action="store_true")
parser.add_argument("-p", help='plot grid of order n', action="store_true")
parser.add_argument("-pp", help='plot primes up to order n ', action="store_true")
parser.add_argument("-pnp", help='plot nonprimes up to order n ', action="store_true")
args = parser.parse_args()

# Do it.

mesh = Eisenstein_mesh(abs(args.n))
	
if args.d:
	print 'eisenstein: \tnumber of points generated : \t', mesh.index 
	for e in mesh.points:
		print  e.float_coord

if args.ds:
	print 'eisenstein: \tnumber of points generated : \t', mesh.index 
	for e in mesh.points:
		print  e.exact_coord

if args.de:
	print 'eisenstein: \tnumber of points generated : \t', mesh.index 
	for e in mesh.points:
		print  'eisenstein:\t',e.eisenstein_coord, '\t\tprime:\t ', e.prime

if args.da:
	print 'eisenstein: \tnumber of points generated : \t', mesh.index 
	for e in mesh.points:
		print e.float_coord,'\t\t', e.exact_coord,',\t', e.eisenstein_coord,'\t', 'prime: ',e.prime

if args.p:
    print 'eisenstein: \tnumber of points generated : \t', mesh.index 
    complex_plot(grid[0])

if args.pp:
    primes = []
    for e in mesh.points:
	if e.prime:
		primes.append(e.float_coord)
    complex_plot(primes)

if args.pnp:
    print 'eisenstein: \tnumber of points generated : \t', mesh.index 
    nonprimes = []
    for e in mesh.points:
	if not e.prime:
		nonprimes.append(e.float_coord)
    print 'eisenstein non primes plotted : ', len(nonprimes)
    complex_plot(nonprimes)
